use calamine::{open_workbook_auto, Reader};
use aes::Aes256;
use cbc::{Decryptor, cipher::{block_padding::Pkcs7, BlockDecryptMut, KeyIvInit}};
use base64::{Engine, engine::general_purpose};
use csv;

// AES-256-CBC 类型别名
pub type Aes256CbcDec = Decryptor<Aes256>;

/// 读取 CSV 文件，解密指定列，并导出到新 CSV 文件
/// 
/// # 参数
/// - input_path: 输入 CSV 路径
/// - output_path: 输出 CSV 路径
/// - decrypt_columns: 需解密的列名（如["A", "C"]）
/// - show_columns: 需展示的列名（如["A", "B", "C"]）
/// - key: 32字节密钥
/// - iv: 16字节初始向量
pub fn process_csv_with_decrypt(
    input_path: &str,
    output_path: &str,
    decrypt_columns: &[&str],
    show_columns: &[&str],
    key: &[u8; 32],
    iv: &[u8; 16],
) -> Result<(), Box<dyn std::error::Error>> {
    // 读取CSV文件
    let mut rdr = csv::Reader::from_path(input_path)?;
    
    // 获取表头
    let headers = rdr.headers()?.clone();
    let header_vec: Vec<String> = headers.iter().map(|h| h.to_string()).collect();

    // 需要展示的列索引
    let show_indices: Vec<usize> = show_columns.iter()
        .map(|col| header_vec.iter().position(|h| h == col).ok_or("Show column not found"))
        .collect::<Result<_, _>>()?;
    
    // 需要解密的列索引
    let decrypt_indices: Vec<usize> = decrypt_columns.iter()
        .map(|col| header_vec.iter().position(|h| h == col).ok_or("Decrypt column not found"))
        .collect::<Result<_, _>>()?;

    // 创建输出CSV写入器
    let mut wtr = csv::Writer::from_path(output_path)?;
    
    // 写入表头
    wtr.write_record(show_columns)?;

    // 处理数据行
    for result in rdr.records() {
        let record = result?;
        let mut new_row = Vec::new();
        
        for &idx in &show_indices {
            let mut cell_val = record.get(idx).unwrap_or("").to_string();
            
            if decrypt_indices.contains(&idx) {
                // base64 解码 + AES 解密
                if let Ok(ciphertext) = general_purpose::STANDARD.decode(&cell_val) {
                    let cipher = Aes256CbcDec::new(key.into(), iv.into());
                    let mut buf = ciphertext.clone();
                    if let Ok(plaintext) = cipher.decrypt_padded_mut::<Pkcs7>(&mut buf) {
                        cell_val = String::from_utf8_lossy(plaintext).to_string();
                    }
                }
            }
            new_row.push(cell_val);
        }
        wtr.write_record(new_row)?;
    }
    
    wtr.flush()?;
    Ok(())
}

/// 读取 Excel 文件，解密指定列，并导出到新 Excel 文件
/// 
/// # 参数
/// - input_path: 输入 Excel 路径
/// - output_path: 输出 Excel 路径
/// - decrypt_columns: 需解密的列名（如["A", "C"]）
/// - show_columns: 需展示的列名（如["A", "B", "C"]）
/// - key: 32字节密钥
/// - iv: 16字节初始向量
pub fn process_excel_with_decrypt(
    input_path: &str,
    output_path: &str,
    decrypt_columns: &[&str],
    show_columns: &[&str],
    key: &[u8; 32],
    iv: &[u8; 16],
) -> Result<(), Box<dyn std::error::Error>> {
    // 检查文件扩展名，决定使用哪种处理方式
    if input_path.ends_with(".csv") {
        return process_csv_with_decrypt(input_path, output_path, decrypt_columns, show_columns, key, iv);
    }
    
    // 打开 Excel 文件
    let mut workbook = open_workbook_auto(input_path)?;
    let range = workbook.worksheet_range_at(0).ok_or("No sheet found")??;

    // 获取表头
    let headers: Vec<String> = range.rows().next()
        .ok_or("No rows found")?
        .iter()
        .map(|cell| cell.to_string())
        .collect();

    // 需要展示的列索引
    let show_indices: Vec<usize> = show_columns.iter()
        .map(|col| headers.iter().position(|h| h == col).ok_or("Show column not found"))
        .collect::<Result<_, _>>()?;
    // 需要解密的列索引
    let decrypt_indices: Vec<usize> = decrypt_columns.iter()
        .map(|col| headers.iter().position(|h| h == col).ok_or("Decrypt column not found"))
        .collect::<Result<_, _>>()?;

    // 收集新表内容
    let mut new_rows = Vec::new();
    // 添加表头
    new_rows.push(show_columns.iter().map(|s| s.to_string()).collect::<Vec<_>>());

    for row in range.rows().skip(1) {
        let mut new_row = Vec::new();
        for &idx in &show_indices {
            let mut cell_val = row.get(idx).map(|c| c.to_string()).unwrap_or_default();
            if decrypt_indices.contains(&idx) {
                // base64 解码 + AES 解密
                if let Ok(ciphertext) = general_purpose::STANDARD.decode(&cell_val) {
                    let cipher = Aes256CbcDec::new(key.into(), iv.into());
                    let mut buf = ciphertext.clone();
                    if let Ok(plaintext) = cipher.decrypt_padded_mut::<Pkcs7>(&mut buf) {
                        cell_val = String::from_utf8_lossy(plaintext).to_string();
                    }
                }
            }
            new_row.push(cell_val);
        }
        new_rows.push(new_row);
    }

    // 写入新 Excel（csv 格式，简单实现）
    let mut wtr = csv::Writer::from_path(output_path)?;
    for row in new_rows {
        wtr.write_record(row)?;
    }
    wtr.flush()?;
    Ok(())
} 